#include "SystemData.h"
#include "BallConstraint.h"
#include "MultibodySystem.h"
#include "CoordinateProjection.h"
#include "Utils.h"
#include "SystemAcceleration.h"

#include <iostream>

#define max(x,y) ( (x) < (y) ? (y) : (x) )
#define min(x,y) ( (x) < (y) ? (x) : (y) )

using namespace RSIM;
using namespace std;

///////////////////////////////////////////////////////

void SystemData::setCoordinateProjectionParameters(const double& Tol, const int& niter){
        Utils::rsim_assert(cpr_,
                           "SystemData not finalized (call SystemData::finalize()) or no constraints added to the system");
        cpr_->setConstraintTolerance(Tol);
        cpr_->setMaxIterations(niter);
}

///////////////////////////////////////////////////////

SystemData::SystemData():SystemDataUtils(),isFinalized_(false),nSys_(0),cpr_(NULL){
}

///////////////////////////////////////////////////////

void SystemData::addGravityLikeAcceleration(const Vect3& accl){
        Gravity_ = new SystemAcceleration(accl);
        for(int i=0;i<nSys_;++i)
                system_[i]->addGravityLikeAcceleration(Gravity_);
}

///////////////////////////////////////////////////////

void SystemData::addSystemData(MultibodySystem *Sys, State *state){
        if(!isFinalized_){
                Sys->wgetSysId() = nSys_++;
                system_.addNode(Sys);
                state_.addNode(state);
                state->initTime(&time_);
        }
        else{
                cout<<"\nCannot add System Data after the function SystemData::finalize() is called\n";
                exit(0);
        }
}

///////////////////////////////////////////////////////

void SystemData::extractYFromState(Vector& Y)const{
        int tmp_var = 0;
        double *ywptr = Y.wPtr();
        const int *ny = NY_.Ptr();
        for(int i=0;i<nSys_;++i){
                state_[i]->getY(ywptr+tmp_var);
                tmp_var += ny[i];
        }
}

///////////////////////////////////////////////////////

void SystemData::setY(const Vector& Y){
        int tmp_var = 0;
        const double *yptr = Y.Ptr();
        const int *ny = NY_.Ptr();
        for(int i=0;i<nSys_;++i){
                state_[i]->setY(yptr+tmp_var);
                tmp_var += ny[i];
        }
}

///////////////////////////////////////////////////////

void SystemData::extractYDotFromState(Vector& YDot)const{
        int tmp_var = 0;
        double *ywptr = YDot.wPtr();
        const int *ny = NY_.Ptr();
        for(int i=0;i<nSys_;++i){
                state_[i]->getYDot(ywptr+tmp_var);
                tmp_var += ny[i];
        }
}

///////////////////////////////////////////////////////

Vector SystemData::getUDotAsVector()const{
        Vector op(nY_);
        double *y = op.wPtr();
        int tmp_var = 0;
        const int *ny = NY_.Ptr();
        for(int i=0;i<nSys_;++i){
                state_[i]->getYDot(y+tmp_var);
                tmp_var += ny[i];
        }
        return op;
}

///////////////////////////////////////////////////////

void SystemData::finalize(){
        if(isFinalized_)
                return;
        isFinalized_ = true;
        
        // Extract the number of Y's in operation
        NY_.resize(nSys_);        
        int *ny = NY_.wPtr();
        nY_ = 0;
        for(int i=0;i<nSys_;++i){                
                ny[i] = state_[i]->getNY();
                nY_ += ny[i];
        }
                
        
        // Initialize list iterators        
        system_.finalize();
        state_.finalize();
        for(int i=0;i<nSys_;++i)
                state_[i]->addSysData(this);
        
        constraints_.finalize();
        
        // Search for coupled algebraic constraints
        initAlgebraicConstraints();        
        
}

///////////////////////////////////////////////////////

void SystemData::calcSystemAcceleration(){        
        for(int i=0;i<nSys_;++i)
                system_[i]->calcSystemAcceleration(*state_[i]);
}

///////////////////////////////////////////////////////

void SystemData::constrain(	const char *str, 		
                                const Transform& Parent_T_PM, 	
                                const SystemIndex& ParentSysId, 	
                                const NodeIndex& ParentSysNodeId, 	
                                const Transform& Child_T_BF,	
                                const SystemIndex& ChildSysId, 	
                                const NodeIndex& ChildSysNodeId
                        ){
        //////////////////////////////////////////////////////////////////
        // Some basic safety checks.
        Utils::rsim_assert(!isFinalized_,
                          "Cannot add constraints after the function SystemData::finalize() is called");
        
        // upper bound of Systems
        Utils::rsim_assert(ParentSysId.id()<nSys_,
                          "SystemData::constrain(): Parent System Id is too high");
                        
        Utils::rsim_assert(ChildSysId.id()<nSys_,
                        "SystemData::constrain(): Child System Id is too high");
                        
        Utils::rsim_assert(ParentSysId.id()>-1,
                        "SystemData::constrain(): Parent system id must be greater than or equal to zero");
        
        Utils::rsim_assert(ChildSysId.id()>-1,
                  "SystemData::constrain(): Child system id must be greater than or equal to zero. To define a constraint between body and ground, use the function SystemData::constrainToGround()");
                        
        // upper bound for number of nodes in each system
        Utils::rsim_assert(ParentSysNodeId.id()<system_[ParentSysId.id()]->getNNodes(), 
                        "SystemData::constrain(): Node Id of the parent is too high");
        
        Utils::rsim_assert(ChildSysNodeId.id()<system_[ChildSysId.id()]->getNNodes(),
                "SystemData::constrain(): Node Id of the child is too high");
        
        // Cannot constraint a given body with itself
        Utils::rsim_assert(!(
                        ParentSysId.id()==ChildSysId.id() && 
                        ParentSysNodeId.id() == ChildSysNodeId.id()
                        ),
                "SystemData::constrain(): Cannot constrain a body with itself"
                );
        // Checks end here.
        //////////////////////////////////////////////////////////////////
        //////////////////////////////////////////////////////////////////
        const int& constraint_id = constraints_.getNNodes();
        if(!strcmp(str,"Ball")){                                                
                
                constraints_.addNode(
                                new BallConstraint(Parent_T_PM, system_[ParentSysId.id()],ParentSysNodeId.id(),
                                                Child_T_BF, system_[ChildSysId.id()], ChildSysNodeId.id(), constraint_id,
                                                   state_[ParentSysId.id()],state_[ChildSysId.id()])
                                        );
        }
        else{
                cout<<"\nInvalid Joint. Possible joint types are:\n <Ball>,\n <Pin>,\n "<<
                "<PrismaticZ>,\n <PrismaticX>,\n <Cylinder>,\n <XYZ>,\n <Free>,\n <Weld>,\n";
                exit(0);
        }        
}

///////////////////////////////////////////////////////

void SystemData::constrainToGround(     const char *str,                
                                        const Transform& Parent_T_PM,   
                                        const SystemIndex& ParentSysId,         
                                        const NodeIndex& ParentSysNodeId,       
                                        const Transform& Child_T_BF
                        ){
        
        Utils::rsim_assert(!isFinalized_,
                          "Cannot add constraints after the function SystemData::finalize() is called");
                          
        // upper bound of Systems
        Utils::rsim_assert(ParentSysId.id()<nSys_,
                        "SystemData::constrain(): Parent System Id is too high");
                                                        
        Utils::rsim_assert(ParentSysId.id()>-1,
                        "SystemData::constrain(): Parent system id must be greater than or equal to zero");
                                        
        // upper bound for number of nodes in each system
        Utils::rsim_assert(ParentSysNodeId.id()<system_[ParentSysId.id()]->getNNodes(), 
                        "SystemData::constrain(): Node Id of the parent is too high");
                                        
        const int& constraint_id = constraints_.getNNodes();
        if(!strcmp(str,"Ball")){
                constraints_.addNode(
                                new BallConstraint(Parent_T_PM, system_[ParentSysId.id()],ParentSysNodeId.id(),
                                                Child_T_BF, NULL, -1,constraint_id,state_[ParentSysId.id()],NULL)
                                                );
        }
        else{
                cout<<"\nInvalid Joint. Possible joint types are:\n <Ball>,\n <Pin>,\n "<<
                "<PrismaticZ>,\n <PrismaticX>,\n <Cylinder>,\n <XYZ>,\n <Free>,\n <Weld>,\n";
                exit(0);
        }
}

///////////////////////////////////////////////////////

void SystemData::updSystemPosition(){
        // ensure that the function finalize() is called before doing anything else.
        
        #ifdef RSIM_SAFE
        Utils::rsim_assert(isFinalized_,
                        "The function SystemData::finalize() must be called before calling SystemData::updSystemPosition()");
        #endif
        
        for(int i=0;i<nSys_;++i){
                // Create new labels for MultibodySystem and State
                MultibodySystem& MBS = *(system_[i]);
                State& state = *(state_[i]);
                
                MBS.updSystemPosition(state);
        }                
}

///////////////////////////////////////////////////////

void SystemData::attemptSatisfyPositionLevelConstraints(){        
        // If algebraic constraints are defined in the system, attempt to satisfy them.
        if(cpr_)                                                
                cpr_->attemptSatisfyPositionLevelConstraint();                        
}

///////////////////////////////////////////////////////

SystemData::~SystemData(){
        const int& nconstraints = constraints_.getLength();
        for(int i=0;i<nconstraints;++i)
                delete constraints_[i];
        delete cpr_;
        delete Gravity_;
}

///////////////////////////////////////////////////////

void SystemData::initAlgebraicConstraints(){
        const int& nConstraints = constraints_.getLength();

        // Nothing to do if there are no constraints in teh system.
        if(nConstraints){
                
                // Create a new object for coordinate projection.
                cpr_ = new CoordinateProjection(ccouple_,
                                                system_,
                                                constraints_,
                                                state_);

                for(int i=0;i<nConstraints;++i){                         
                        
                        constraints_[i]->init();
                        
                        // Compute the CommonParentId for each constraint and set the path of the constrained
                        // body relative to the common parent
                        const int& common_parent = this->calcCommonParentForConstraint(system_,constraints_[i]);
                        constraints_[i]->setCommonParentId(common_parent);
                        
                        // Calculate the body FM transforms for both the bodies involved in each constraint
                        constraints_[i]->initConstraintBodyFixedTransforms();
                }                                                
                // Next to check for coupling between constraints. This is simply done to 
                // exploit the block structure of the system-wide constraint matrix.

                // For a single constraint there is no question of coupling
                if(nConstraints==1){
                        List <int> tmp;
                        tmp.addNode(0);
                        ccouple_.addNode(tmp);
                }
                else{                
                        List <List<int> > tmp_list;

                        for(int i=0;i<nConstraints-1;++i){
                                for(int j=i+1;j<nConstraints;++j){
                                        // check if constraints i and j are coupled
                                        const bool& bool_var = areConstraintsCoupled(constraints_[i],
                                                                                     constraints_[j]);
                                        if(bool_var){
                                                List <int> tmp_list_1;
                                                tmp_list_1.addNode(i);
                                                tmp_list_1.addNode(j);
                                                tmp_list.addNode(tmp_list_1);
                                                tmp_list_1.clear();
                                        }
                                }
                        }
                                                                                        
                        // This figures out the coupling within different constraints. This is different form 
                        // the loop above.
                        this->sortCoupledConstraintList_1(nConstraints,tmp_list,ccouple_);
                }                                
                cpr_->init();                
        }//if(nConstraints)
}

///////////////////////////////////////////////////////